package cn.geminis.crypto.csp.soft.rsa;

import cn.geminis.crypto.core.key.PrivateKey;
import cn.geminis.crypto.core.key.PublicKey;
import cn.geminis.crypto.csp.AbstractAgreement;
import cn.geminis.crypto.csp.parameter.CalcAgreementCipherParameters;
import cn.geminis.crypto.csp.parameter.InitAgreementCipherParameters;
import org.bouncycastle.asn1.ASN1ObjectIdentifier;
import org.bouncycastle.asn1.ASN1Sequence;
import org.bouncycastle.asn1.DERSequence;
import org.bouncycastle.asn1.DERSequenceGenerator;
import org.bouncycastle.asn1.pkcs.PrivateKeyInfo;
import org.bouncycastle.asn1.x509.SubjectPublicKeyInfo;
import org.bouncycastle.asn1.x9.ECNamedCurveTable;
import org.bouncycastle.asn1.x9.X9ECParameters;
import org.bouncycastle.crypto.AsymmetricCipherKeyPair;
import org.bouncycastle.crypto.agreement.ECDHBasicAgreement;
import org.bouncycastle.crypto.generators.ECKeyPairGenerator;
import org.bouncycastle.crypto.params.AsymmetricKeyParameter;
import org.bouncycastle.crypto.params.ECKeyGenerationParameters;
import org.bouncycastle.crypto.params.ECNamedDomainParameters;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.math.BigInteger;
import java.security.SecureRandom;

/**
 * @author Allen
 */
public class ECDHAgreement extends AbstractAgreement {

    private ECDHBasicAgreement ecdhBasicAgreement = new ECDHBasicAgreement();
    private PrivateKey privateKey;
    private AsymmetricKeyParameter tempPrivateKey;
    private AsymmetricKeyParameter tempPublicKey;

    public ECDHAgreement(PublicKey publicKey, PrivateKey privateKey) {
        super(publicKey);
        this.privateKey = privateKey;

        var name = "P-256";
        var oid = ECNamedCurveTable.getOID(name);
        var ecSpec = ECNamedCurveTable.getByName(name);
        var domainParameters = new ECNamedDomainParameters(oid, ecSpec.getCurve(), ecSpec.getG(), ecSpec.getN());
        var generationParameters = new ECKeyGenerationParameters(domainParameters, new SecureRandom());
        var keyPairGenerator = new ECKeyPairGenerator();
        keyPairGenerator.init(generationParameters);

        var keyPair = keyPairGenerator.generateKeyPair();
        this.tempPrivateKey = keyPair.getPrivate();
        this.tempPublicKey = keyPair.getPublic();
    }

    @Override
    public void init(InitAgreementCipherParameters param) {
        this.ecdhBasicAgreement.init(this.tempPrivateKey);
    }

    @Override
    public BigInteger calculateAgreement(CalcAgreementCipherParameters param) {
        return ecdhBasicAgreement.calculateAgreement(param.getTempPublicKey().getKeyParameter());
    }

    @Override
    public byte[] getSession() {
        try (var stream = new ByteArrayOutputStream()) {
            DERSequenceGenerator generator = new DERSequenceGenerator(stream);
            generator.addObject(new PublicKey(this.tempPublicKey).getSubjectPublicKeyInfo());
            generator.addObject(new PrivateKey(this.tempPrivateKey).getPrivateKeyInfo());
            generator.close();
            return stream.toByteArray();
        } catch (IOException e) {
            throw new RuntimeException("获取密钥协商器Session错误", e);
        }
    }

    @Override
    public void setSession(byte[] session) {
        try {
            var sequence = (ASN1Sequence) DERSequence.fromByteArray(session);
            var publicKeyInfo = SubjectPublicKeyInfo.getInstance(sequence.getObjectAt(0));
            var privateKeyInfo = PrivateKeyInfo.getInstance(sequence.getObjectAt(1));
            this.tempPublicKey = new PublicKey(publicKeyInfo).getKeyParameter();
            this.tempPrivateKey = new PrivateKey(privateKeyInfo).getKeyParameter();
        } catch (IOException e) {
            throw new RuntimeException("设置密钥协商器Session错误", e);
        }
    }

    @Override
    public PublicKey getTempPublicKey() {
        return new PublicKey(this.tempPublicKey);
    }

    @Override
    public int getFieldSize() {
        return 64;
    }

    @Override
    public void close() {

    }
}
